#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/7/28 3:04 下午
# @Author  : WangZhixing

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np

def visualize(embedding, value, edge_index=None):
    '''
    对embedding进行作图，判断分类器效果如何。
    :param edge_index: 如果画边，则需要edge_index[2,edge_num]
    :param value: 每个node的label类
    :param embedding: [[],[]]格式，表示每个类的embedding向量
    '''
    # # TSNE提供了一种有效的降维方式，让我们对高于2维数据的聚类结果以二维的方式展示出来

    if type(embedding) is not np.ndarray:
        z = embedding.detach().cpu().numpy()
    else:
        z = embedding
    if z.shape[1] is not 2:
        z = TSNE(n_components=2).fit_transform(embedding.detach().cpu().numpy())

    plt.figure(figsize=(16, 12))
    plt.xticks([])
    plt.yticks([])
    plt.scatter(z[:, 0], z[:, 1], s=30, c=value, cmap="Set2", zorder=2)
    if edge_index is not None:
        for edge in edge_index.t():
            start, end = edge
            plt.plot([z[start][0], z[end][0]], [z[start][1], z[end][1]], color='gray', alpha=0.1, zorder=1)
    plt.show()
